import sys
import os

# 将项目根目录添加到Python路径中，以便导入dragonquant
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import dragonquant as dq
from dragonquant.api import order, log
from dragonquant.indicators import calculate_sma

class DualMovingAverage(dq.Strategy):
    """
    双均线策略
    - 当短期均线上穿长期均线时，买入
    - 当短期均线下穿长期均线时，卖出
    """
    def init(self):
        # 1. 设置策略参数
        self.stock = '000001.SH'  # 交易标的：上证指数
        self.short_window = 5      # 短期均线周期
        self.long_window = 20      # 长期均线周期
        self.bought = False  # 记录是否已买入
        log("双均线策略初始化完成。")

    def next(self, context):
        # 2. 获取历史数据
        hist = context.data.history(self.stock, ['close'], self.long_window + 1, context.dt)
        if hist.empty or len(hist) < self.long_window + 1:
            return
        
        log(f"在 {context.dt} 执行策略...")
        
        # 3. 计算均线
        ma_short = calculate_sma(hist['close'], self.short_window)
        ma_long = calculate_sma(hist['close'], self.long_window)
        
        # 4. 信号生成与交易
        current_price = hist['close'].iloc[-1]
        
        # 金叉：短期均线上穿长期均线
        if ma_short > ma_long and not self.bought:
            # 买入逻辑
            shares_to_buy = int(context.portfolio.cash / current_price // 100) * 100  # 买入100的整数倍
            if shares_to_buy > 0:
                order(self.stock, shares_to_buy)
                self.bought = True
                log(f"买入 {self.stock}，数量：{shares_to_buy}，价格：{current_price:.2f}")
        # 死叉：短期均线下穿长期均线
        elif ma_short < ma_long and self.bought:
            # 卖出逻辑（全仓卖出）
            order(self.stock, 0, style='target')
            self.bought = False
            log(f"卖出 {self.stock}，价格：{current_price:.2f}")

if __name__ == '__main__':
    # 1. 初始化回测引擎
    cerebro = dq.Cerebro()

    # 2. 添加策略
    cerebro.add_strategy(DualMovingAverage)

    # 3. 设置股票池
    cerebro.set_universe(['000001.SH'])

    # 4. 设置回测参数
    cerebro.configure(
        start_date='2022-01-01',
        end_date='2023-12-31',
        initial_cash=1000000,
        data_source='fake'
    )

    # 5. 运行回测
    results = cerebro.run()

    # 6. 分析和可视化结果
    if results:
        results.generate_report()
        results.plot_equity_curve()
